#!/usr/bin/env python
# -*- coding: utf-8 -*-

import os
import sys
import io
import json

from PIL import Image
from PIL import ImagePalette

import numpy as np
import struct

from loguru import logger

sys.path.append(os.getcwd())
from models.data_model import Acquisition
from utils.app_config import AppConfig

### 将matlab生成的图像结构转换为PIL Image进行存储和计算
class PAIAnalysis(object):
    def __new__(cls, *args, **kwargs):
        if not hasattr(cls, '_instance'):
            cls._instance = super().__new__(cls, *args, **kwargs)
        return cls._instance

    def __init__(self):
        super().__init__()

        self.current_acquisition_id = 0
        self.current_acquisition_path = None
        self.current_acquisition_file = None
        self.paus_image = {}

        crop_region = AppConfig().getCropROIRegion()
        self.crop_box = (int(crop_region[0]), int(crop_region[1]), int(crop_region[2]), int(crop_region[3]))

        logger.debug("create an instance of PAIAnalysis successfully.")

    ## get recon methods list
    def getsupportedReconMethods(self):
        recon_methods_list = AppConfig().getReconstructionMethods()
        return recon_methods_list

    def setCurrentAcquisitionID(self, acquisition_id):
        self.current_acquisition_id = acquisition_id
    
    def getCurrentAcquisitionID(self):
        return self.current_acquisition_id
    
    def getCurrentAcquisitionPath(self):
        existing_acq = Acquisition.query.filter(Acquisition.id == self.current_acquisition_id).one_or_none()
        if existing_acq is None:
                return None
        
        self.current_acquisition_path = existing_acq.acq_data_path
        return self.current_acquisition_path
    
    def setCurrentAcquisitionFile(self, bin_file):
        self.current_acquisition_file = bin_file

    def getCurrentAcquisitionFile(self):
        return self.current_acquisition_file
    
    def getCropbox(self):
        return self.crop_box

    ##  acq_id:  acquisition id
    ##  modality: "US" or "PA"
    ##  bmp_data: bmp格式的图像
    ## 在内部存为PIL Image，外部需要时进行转换
    def setImage(self, acq_id, modality, bmp_data, *params):
        image = self.__generatePILImage(bmp_data, params)
        if self.paus_image.get(acq_id):
            if modality == "US":
                self.paus_image[acq_id].update({"us_image":image})
            elif modality == "PA":
                self.paus_image[acq_id].update({"pa_image":image})
            else:
                logger.error("input parameter modality is invalid. only PA and US are supported!")
        else:
            if modality == "US":
                self.paus_image[acq_id] = {"us_image":image}
            elif modality == "PA":
                self.paus_image[acq_id] = {"pa_image":image}
            else:
                logger.error("input parameter modality is invalid. only PA and US are supported!")

    ##  acq_id:  acquisition id
    ##  modality: "US" or "PA"
    ##  bmp_data: bmp格式的图像
    def setCropImage(self, acq_id, modality, bmp_data, *params):
        image = self.__generatePILImage(bmp_data, params)
        if self.paus_image.get(acq_id):
            if modality == "US":
                self.paus_image[acq_id].update({"us_crop_image":image})
            elif modality == "PA":
                self.paus_image[acq_id].update({"pa_crop_image":image})
            else:
                logger.error("input parameter modality is invalid. only PA and US are supported!")
        else:
            if modality == "US":
                self.paus_image[acq_id] = {"us_crop_image":image}
            elif modality == "PA":
                self.paus_image[acq_id] = {"pa_crop_image":image}
            else:
                logger.error("input parameter modality is invalid. only PA and US are supported!")
    
    def getImageList(self, modality, acq_id_list):
        image_list = []
        if modality == "US":
            for acq_id in self.paus_image:
                if acq_id in acq_id_list:
                    if self.paus_image[acq_id].get("us_image"):
                        image_list.append(self.paus_image[acq_id].get("us_image")) 
        elif modality == "PA":
            for acq_id in self.paus_image:
                if acq_id in acq_id_list:
                    if self.paus_image[acq_id].get("pa_image"):
                        image_list.append(self.paus_image[acq_id].get("pa_image"))
        
        return image_list

    
    def getCropImageList(self, modality, acq_id_list):
        image_list = []
        if modality == "US":
            for acq_id in self.paus_image:
                if acq_id in acq_id_list:
                    if self.paus_image[acq_id].get("us_crop_image"):
                        image_list.append(self.paus_image[acq_id].get("us_crop_image")) 
        elif modality == "PA":
            for acq_id in self.paus_image:
                if acq_id in acq_id_list:
                    if self.paus_image[acq_id].get("pa_crop_image"):
                        image_list.append(self.paus_image[acq_id].get("pa_crop_image"))
        
        return image_list
    
    def extractImagePixelBytes(self, image):
        pixel_matrix = np.array(image, 'uint8')
        array_1d = pixel_matrix.flatten("F")
        # 将获取数据转换为bytes
        pixel_bytes = array_1d.tobytes()
        return pixel_bytes

    def __generatePILImage(self, bmp_data, *params):
        image_object = io.BytesIO()
        image_object.write(bmp_data)
        image = Image.open(image_object)
        #print(params)
        if len(params) == 1 and isinstance(params[0], tuple):
            if len(params[0]) == 1 and isinstance(params[0][0], dict):
                #print(params[0][0])
                image.info.update(params[0][0])
        #image.show()
        #print(image.info)
        return image
    
    def image2BMPBytes(self, image):
        image_object = io.BytesIO()
        image.save(image_object, "bmp")
        bmp_data_bytes = image_object.getvalue()
        return bmp_data_bytes

    ## 裁剪图像，并转换为bmp bytes，输入为bmp字节流和crop box，输出为bmp字节流
    def cropImage(self, bmp_data, crop_box):
        if len(crop_box) != 4:
            logger.error("Errors happen in cropImage because the crop_region is invalid.")
            return None
        image = self.__generatePILImage(bmp_data)
        crop_image = image.crop((int(crop_box[0]), int(crop_box[1]), int(crop_box[2]), int(crop_box[3])))
        #crop_image.show()
        #crop_bmp_data = self.__image2BMPBytes(crop_image)
        return crop_image
    
    ## 将图像的感兴趣区域提取为包括感兴趣区域的最小范围图像。感兴趣区外填充为0.
    ## roi_type 目前仅支持rectangle。
    def extractROIImage(self, image, roi_type, point_list):
        if roi_type == "rectangle":
            roi_rectangle = image.crop((int(point_list[0]), int(point_list[1]), int(point_list[2]), int(point_list[3])))
            return roi_rectangle
        if roi_type == "polygon":
            ### TBD
            return None

    ## Fusion the US image and PA image
    def pausFusion(self, acq_id, blending_factor):
        if not self.paus_image.get(acq_id):
            return None
        pa_image = self.paus_image[acq_id]["pa_image"]
        #print(pa_image.info)
        us_image = self.paus_image[acq_id]["us_image"]
        us_gray_image = us_image.convert("L")
        us_rgb_image = us_gray_image.convert("RGB")
        #us_rgb_image.show()
        pa_rgb_image = pa_image.convert("RGB")
        #print(pa_rgb_image.info)
        #pa_rgb_image.show()

        fusion_image = Image.blend(us_rgb_image, pa_rgb_image, blending_factor)
        #print(fusion_image.info)
        #fusion_image.show()

        #fusion_bmp_data = self.image2BMPBytes(fusion_image)
        # with open('pixeldata.dat', 'wb') as f:
        #     f.write(bmp_data)
        # self.showImage(bmp_data)

        return fusion_image
    
    ## method for testing
    def showImage(self, image_data):
        image_object = io.BytesIO()
        image_object.write(image_data)
        image_show = Image.open(image_object)
        image_show.show()

    ## generate BMP data with palette
    def generateBMPWithPaletteImage(self, data_matrix, dimension, palette):
        width = int(dimension[0])
        height = int(dimension[1])
        # 定义bmp文件头
        bmp_header = struct.pack('<2sIHHI', b'BM', 54 + width * height + 256 * 4, 0, 0, 54 + 256 * 4)

        # 定义bmp信息头
        bmp_info_header = struct.pack('<IIIHHIIIIII', 40, width, height, 1, 8, 0, width * height, 0, 0, 256, 0)
        # 将数组矩阵转换为像素数据和调色板数据
        pixel_data = bytes(data_matrix[::-1].reshape(width * height))
        palette_data = bytes(palette.reshape(256 * 4))

        # 创建一个BytesIO对象
        output = io.BytesIO()
        output.write(bmp_header)
        output.write(bmp_info_header)
        output.write(palette_data)
        output.write(pixel_data)

        # 将BytesIO对象中的数据保存到变量中
        bmp_data = output.getvalue()
        return bmp_data

    ## generate BMP Gray data
    def generateBMPGrayImage(self, data_matrix, dimension):
        width = int(dimension[0])
        height = int(dimension[1])
        # 定义bmp文件头
        bmp_header = struct.pack('<2sIHHI', b'BM', 54 + width * height, 0, 0, 54)

        # 定义bmp信息头
        bmp_info_header = struct.pack('<IIIHHIIIIII', 40, width, height, 1, 8, 0, width * height, 0, 0, 0, 0)
        # 将数组矩阵转换为像素数据和调色板数据
        pixel_data = bytes(data_matrix[::-1].reshape(width * height))
        #pixel_data = bytes(data_matrix.reshape(width * height))

        # 创建一个BytesIO对象
        output = io.BytesIO()
        output.write(bmp_header)
        output.write(bmp_info_header)
        output.write(pixel_data)

        # 将BytesIO对象中的数据保存到变量中
        bmp_data = output.getvalue()
        return bmp_data

    ## generate BMP data with RBG or RGBA
    def generateBMPImage(self, data_matrix, dimension):
        print(f"data_matrix shape: ", data_matrix.shape)
        print(f"data_matrix dimension: ", dimension)
        width = int(dimension[0])
        height = int(dimension[1])
        channels = int(dimension[2])
        # 定义bmp文件头
        bmp_header = struct.pack('<2sIHHI', b'BM', 54 + width * height * channels, 0, 0, 54)

        # 定义bmp信息头
        bmp_info_header = struct.pack('<IIIHHIIIIII', 40, width, height, 1, 8 * channels, 0, width * height * channels, 0, 0, 0, 0)
        # 将数组矩阵转换为像素数据和调色板数据
        #image = Image.fromarray(data_matrix.reshape(channels, height, width))
        matlab_matrix = data_matrix.reshape(channels, height, width)
        # print(f"matlab_matrix after data_matrix reshape: ", matlab_matrix.shape)
        matlab_matrix = matlab_matrix.transpose(2,1,0)
        # print(f"matlab_matrix after 1st transpose: ", matlab_matrix.shape)
        matlab_matrix = matlab_matrix.transpose(1,0,2)
        # print(f"matlab_matrix after 2nd transpose: ", matlab_matrix.shape)
        # matlab_image = Image.fromarray(matlab_matrix)
        # matlab_image.show()

        #pixel_data = bytes(data_matrix[::-1].reshape(width * height * channels))
        pixel_data = bytes(matlab_matrix[:, :, ::-1])

        # 创建一个BytesIO对象
        output = io.BytesIO()
        output.write(bmp_header)
        output.write(bmp_info_header)
        output.write(pixel_data)

        # 将BytesIO对象中的数据保存到变量中
        bmp_data = output.getvalue()
        return bmp_data

class MatlabMatEncoder(json.JSONEncoder):
    def default(self, obj):
        if isinstance(obj, np.ndarray):
            return obj.tolist()
        return super().default(obj)

#################################################################
if __name__ == '__main__':
    # define a global single variable
    pai_analysis = PAIAnalysis()
    # # 创建一个随机的8位数组矩阵
    #matrix = np.random.randint(0, 256, size = (height, width, 3), dtype=np.uint8)
    # dimension = np.ndarray([256, 256])

    # # 创建调色板数据
    # palette = np.zeros((256, 4), dtype=np.uint8)
    # palette[:, 0] = np.arange(256, dtype=np.uint8)
    # palette[:, 1] = np.arange(256, dtype=np.uint8)
    # palette[:, 2] = np.arange(256, dtype=np.uint8)
    # palette[:, 3] = 0
    random_matrix = np.random.randint(0, 256, size = (2, 3, 4), dtype=np.uint8)
    print(random_matrix)

    image = Image.open("D:/Project/pabone/matlab_scripts/image_signal_analysis/mineral_result.bmp")
    print(f"Image size:", image.size)
    array = np.array(image)
    print(f"Image array shape:", array.shape) 

    dimension = (image.size[0], image.size[1], 3)
    bmp_data = pai_analysis.generateBMPImage(array, dimension)
    with open('output.data', 'wb') as f:
         f.write(bmp_data)

